fix(masque): avoid polling completed relay task handles - #128

Benchmark Results

Performance Comparison

Summary

Configuration

• Regression threshold: >10% slower
• Improvement threshold: >10% faster
• Measurements: Mean execution time

Review

The bug is real and correctly diagnosed. In the stream forwarding loop, the tokio::select! polls &mut reader_handle / &mut writer_handle. When one wins, its JoinHandle output is consumed; the old cleanup then unconditionally did reader_handle.await / writer_handle.await, re-polling the already-completed handle and panicking with Tokio's JoinHandle polled after completion. Under panic = "abort" consumers (e.g. ant-node) this takes the whole process down during relay teardown, including upgrade-triggered graceful shutdown.

The fix is correct. The core invariant holds: in tokio::select!, a branch future's output is consumed iff that branch is selected as the winner (select! short-circuits on the first Ready). completed_task is set exactly inside each winning branch, so it precisely tracks which handle was consumed:

• Reader/Writer won → the completed handle is dropped (task already finished, socket clone released), the other handle is aborted + awaited.
• Neither (handover cancel, or the inline dir2 stream→UDP branch) → both handles aborted + awaited, matching the original behavior.

The Blocker 1 guarantee is preserved: the reader task owns the moved socket Arc clone; in every path it either completes naturally or is aborted+awaited before finish_forwarding/close_session can lease the socket. The inline dir2 future holding socket2 is dropped automatically when select! returns.

Tests are valid regression coverage — each exercises the exact drop-vs-await decision and would panic if the completed handle were re-polled.

No blocking concerns. The change is minimal, correct, and well-reasoned.

One minor, non-blocking suggestion: the tests cover the extracted drain_forwarding_tasks helper in isolation, not the select! → completed_task wiring itself. An integration test that drives the real forwarding loop through a reader-wins path would guard against a future refactor reintroducing the double-poll — an enhancement, not a defect.

LGTM. 👍

Assessed the integration-test suggestion. I agree with the coverage boundary, but I do not think expanding this PR to construct a live MASQUE session/QUIC stream harness is proportionate. The regression is in the cleanup decision itself, which is exhaustively covered for Reader, Writer and Neither; the production wiring consists of branch-local assignments immediately before handling each irrefutable JoinHandle result. A full forwarding-loop test would require substantially more transport/session setup and would be more brittle than the code under test. The structurally stricter Option<JoinHandle>::take() alternative also requires a broader control-flow refactor here. I am therefore keeping this fix narrow; a future end-to-end relay teardown test would still be useful as separate test-harness work.